import numpy as np
import pandas as pd
from typing import List

PROBABILITY_LIST = [0.1, 0.3, 0.6] # 每个产品对应的概率 总和为1
NUM_EXPERIMENT = 100 # 取样总次数
OUT_SIZE = 10  # 输出次数

class Multi(object):
    """
    multinomial 多项式分布中提取样本
    """
    def __init__(self, probability: List[float], num_experiment: int, out_size: int) -> None:
        self.probability = probability
        rng = np.random.default_rng()
        self.rvs = rng.multinomial(num_experiment, probability, out_size)

    def to_frame(self) -> pd.DataFrame:
        df = pd.DataFrame(
            self.rvs,
            columns=self.probability,
            index=range(1, self.rvs.shape[0] + 1)
        )
        df['total'] = df.sum(axis=1)
        return df.T

class Bi(object):
    """
    binomial 二项分布中提取样本
    """
    def __init__(self, probability: List[float], num_experiment: int, out_size: int) -> None:
        self.out_size = out_size
        rng = np.random.default_rng()
        self.rvs_d = {str(p): rng.binomial(num_experiment, p, out_size) for p in probability}

    def to_frame(self) -> pd.DataFrame:
        df = pd.DataFrame(
            self.rvs_d,
            index=range(1, self.out_size + 1)
        )
        df['total'] = df.sum(axis=1)
        return df.T


if __name__ == '__main__':
    print(Multi(PROBABILITY_LIST, NUM_EXPERIMENT, OUT_SIZE).to_frame())
    print(Bi(PROBABILITY_LIST, NUM_EXPERIMENT, OUT_SIZE).to_frame())
